import numpy as np
import tensorflow as tf

import os
from random import randint
from PIL import Image
from django.conf import settings


class Captcha(object):
    def __init__(self):

        # 图像大小
        self.IMAGE_HEIGHT = 60
        self.IMAGE_WIDTH = 160
        self.MAX_CAPTCHA = 4
        # print("验证码文本最长字符数", self.MAX_CAPTCHA)  # 验证码最长4字符; 我全部固定为4,可以不固定. 如果验证码长度小于4，用'_'补齐
        self.number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
        self.alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's',
                         't', 'u', 'v', 'w', 'x', 'y', 'z']
        self.ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
                         'T', 'U', 'V', 'W', 'X', 'Y', 'Z']

        # 文本转向量
        char_set = self.number + self.alphabet + self.ALPHABET + ['_']  # 如果验证码长度小于4, '_'用来补齐
        self.CHAR_SET_LEN = len(char_set)

        self.X = tf.placeholder(tf.float32, [None, self.IMAGE_HEIGHT * self.IMAGE_WIDTH])
        self.Y = tf.placeholder(tf.float32, [None, self.MAX_CAPTCHA * self.CHAR_SET_LEN])
        self.keep_prob = tf.placeholder(tf.float32)  # dropout
        self.train_phase = tf.placeholder(tf.bool)
        self.sess = tf.Session()
        self.output = self.crack_captcha_cnn()
        self.predict = tf.argmax(tf.reshape(self.output, [-1, self.MAX_CAPTCHA, self.CHAR_SET_LEN]), 2)
        self.saver = tf.train.Saver()
        self.saver.restore(self.sess, tf.train.latest_checkpoint(settings.BASE_DIR + "/unicom/model"))

    # 把彩色图像转为灰度图像（色彩对识别验证码没有什么用）
    def convert2gray(self, img):
        if len(img.shape) > 2:
            gray = np.mean(img, -1)
            # 上面的转法较快，正规转法如下
            # r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
            # gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
            return gray
        else:
            return img

    """
    cnn在图像大小是2的倍数时性能最高, 如果你用的图像大小不是2的倍数，可以在图像边缘补无用像素。
    np.pad(image,((2,3),(2,2)), 'constant', constant_values=(255,))  # 在图像上补2行，下补3行，左补2行，右补2行
    """

    def text2vec(self, text):
        text_len = len(text)
        if text_len > self.MAX_CAPTCHA:
            raise ValueError('验证码最长4个字符')
        vector = np.zeros(self.MAX_CAPTCHA * self.CHAR_SET_LEN)

        def char2pos(c):
            if c == '_':
                k = 62
                return k
            k = ord(c) - 48
            if k > 9:
                k = ord(c) - 55
                if k > 35:
                    k = ord(c) - 61
                    if k > 61:
                        raise ValueError('No Map')
            return k

        for i, c in enumerate(text):
            idx = i * self.CHAR_SET_LEN + char2pos(c)
            vector[idx] = 1
        return vector

    # 向量转回文本
    def vec2text(self, vec):
        char_pos = vec.nonzero()[0]
        text = []
        for i, c in enumerate(char_pos):
            char_at_pos = i  # c/63
            char_idx = c % self.CHAR_SET_LEN
            if char_idx < 10:
                char_code = char_idx + ord('0')
            elif char_idx < 36:
                char_code = char_idx - 10 + ord('A')
            elif char_idx < 62:
                char_code = char_idx - 36 + ord('a')
            elif char_idx == 62:
                char_code = ord('_')
            else:
                raise ValueError('error')
            text.append(chr(char_code))
        return "".join(text)

    """
    #向量（大小MAX_CAPTCHA*CHAR_SET_LEN）用0,1编码 每63个编码一个字符，这样顺利有，字符也有
    vec = text2vec("F5Sd")
    text = vec2text(vec)
    print(text)  # F5Sd
    vec = text2vec("SFd5")
    text = vec2text(vec)
    print(text)  # SFd5
    """
    dataset = []
    labelset = []
    label_map = {}

    # from sklearn.model_selection import train_test_split
    def gen_captcha_text_and_image(self):
        path = 'E:\\captcha\\chinatelecom\\trainset\\train1102\\'
        files = os.listdir(path)
        # print("len(files):",len(files) - 1,)
        random_num = randint(0, len(files) - 1)
        file_name = files[random_num]
        img = Image.open(path + file_name).crop([0, 0, 160, 60]).convert('L')
        (width, height) = img.size
        pixles = img.load()

        for w in range(width):  # 二值化
            for h in range(height):
                if pixles[w, h] > 224:
                    pixles[w, h] = 255
                else:
                    pixles[w, h] = 0
        return file_name[0:4], np.array(img)

    # 生成一个训练batch
    def get_next_batch(self, batch_size=128):
        batch_x = np.zeros([batch_size, self.IMAGE_HEIGHT * self.IMAGE_WIDTH])
        batch_y = np.zeros([batch_size, self.MAX_CAPTCHA * self.CHAR_SET_LEN])

        for i in range(batch_size):
            text, image = self.gen_captcha_text_and_image()
            image = self.convert2gray(image)
            batch_x[i, :] = image.flatten() / 255  # (image.flatten()-128)/128  mean为0
            batch_y[i, :] = self.text2vec(text)
        return batch_x, batch_y

    ####################################################################


    def batch_norm(self, x, beta, gamma, phase_train, scope='bn', decay=0.9, eps=1e-5):
        with tf.variable_scope(scope):
            # beta = tf.get_variable(name='beta', shape=[n_out], initializer=tf.constant_initializer(0.0), trainable=True)
            # gamma = tf.get_variable(name='gamma', shape=[n_out], initializer=tf.random_normal_initializer(1.0, stddev), trainable=True)
            batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], name='moments')
            ema = tf.train.ExponentialMovingAverage(decay=decay)

            def mean_var_with_update():
                ema_apply_op = ema.apply([batch_mean, batch_var])
                with tf.control_dependencies([ema_apply_op]):
                    return tf.identity(batch_mean), tf.identity(batch_var)

            mean, var = tf.cond(phase_train, mean_var_with_update,
                                lambda: (ema.average(batch_mean), ema.average(batch_var)))
            normed = tf.nn.batch_normalization(x, mean, var, beta, gamma, eps)
        return normed

    # 定义CNN
    def crack_captcha_cnn(self, w_alpha=0.01, b_alpha=0.1):
        x = tf.reshape(self.X, shape=[-1, self.IMAGE_HEIGHT, self.IMAGE_WIDTH, 1])

        # 4 conv layer
        w_c1 = tf.Variable(w_alpha * tf.random_normal([3, 3, 1, 32]))
        b_c1 = tf.Variable(b_alpha * tf.random_normal([32]))
        conv1 = tf.nn.bias_add(tf.nn.conv2d(x, w_c1, strides=[1, 1, 1, 1], padding='SAME'), b_c1)
        conv1 = self.batch_norm(conv1, tf.constant(0.0, shape=[32]),
                                tf.random_normal(shape=[32], mean=1.0, stddev=0.02),
                                self.train_phase, scope='bn_1')
        conv1 = tf.nn.relu(conv1)
        conv1 = tf.nn.max_pool(conv1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv1 = tf.nn.dropout(conv1, self.keep_prob)

        w_c2 = tf.Variable(w_alpha * tf.random_normal([3, 3, 32, 64]))
        b_c2 = tf.Variable(b_alpha * tf.random_normal([64]))
        conv2 = tf.nn.bias_add(tf.nn.conv2d(conv1, w_c2, strides=[1, 1, 1, 1], padding='SAME'), b_c2)
        conv2 = self.batch_norm(conv2, tf.constant(0.0, shape=[64]),
                                tf.random_normal(shape=[64], mean=1.0, stddev=0.02),
                                self.train_phase, scope='bn_2')
        conv2 = tf.nn.relu(conv2)
        conv2 = tf.nn.max_pool(conv2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv2 = tf.nn.dropout(conv2, self.keep_prob)

        w_c3 = tf.Variable(w_alpha * tf.random_normal([3, 3, 64, 64]))
        b_c3 = tf.Variable(b_alpha * tf.random_normal([64]))
        conv3 = tf.nn.bias_add(tf.nn.conv2d(conv2, w_c3, strides=[1, 1, 1, 1], padding='SAME'), b_c3)
        conv3 = self.batch_norm(conv3, tf.constant(0.0, shape=[64]),
                                tf.random_normal(shape=[64], mean=1.0, stddev=0.02),
                                self.train_phase, scope='bn_3')
        conv3 = tf.nn.relu(conv3)
        conv3 = tf.nn.max_pool(conv3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv3 = tf.nn.dropout(conv3, self.keep_prob)

        w_c4 = tf.Variable(w_alpha * tf.random_normal([3, 3, 64, 64]))
        b_c4 = tf.Variable(b_alpha * tf.random_normal([64]))
        conv4 = tf.nn.bias_add(tf.nn.conv2d(conv3, w_c4, strides=[1, 1, 1, 1], padding='SAME'), b_c4)
        conv4 = self.batch_norm(conv4, tf.constant(0.0, shape=[64]),
                                tf.random_normal(shape=[64], mean=1.0, stddev=0.02),
                                self.train_phase, scope='bn_4')
        conv4 = tf.nn.relu(conv4)
        conv4 = tf.nn.max_pool(conv4, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
        conv4 = tf.nn.dropout(conv4, self.keep_prob)

        # Fully connected layer
        w_d = tf.Variable(w_alpha * tf.random_normal([2 * 20 * 64, 1024]))
        b_d = tf.Variable(b_alpha * tf.random_normal([1024]))
        dense = tf.reshape(conv4, [-1, w_d.get_shape().as_list()[0]])
        dense = tf.nn.relu(tf.add(tf.matmul(dense, w_d), b_d))
        dense = tf.nn.dropout(dense, self.keep_prob)

        w_out = tf.Variable(w_alpha * tf.random_normal([1024, self.MAX_CAPTCHA * self.CHAR_SET_LEN]))
        b_out = tf.Variable(b_alpha * tf.random_normal([self.MAX_CAPTCHA * self.CHAR_SET_LEN]))
        out = tf.add(tf.matmul(dense, w_out), b_out)
        # out = tf.nn.softmax(out)
        return out

    # 训练
    def train_crack_captcha_cnn(self):
        output = self.crack_captcha_cnn()
        # loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=output, labels=Y))
        loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=output, labels=self.Y))
        # 最后一层用来分类的softmax和sigmoid有什么不同？
        # optimizer 为了加快训练 learning_rate应该开始大，然后慢慢衰
        optimizer = tf.train.AdamOptimizer(learning_rate=0.002).minimize(loss)

        predict = tf.reshape(output, [-1, self.MAX_CAPTCHA, self.CHAR_SET_LEN])
        max_idx_p = tf.argmax(predict, 2)
        max_idx_l = tf.argmax(tf.reshape(self.Y, [-1, self.MAX_CAPTCHA, self.CHAR_SET_LEN]), 2)
        correct_pred = tf.equal(max_idx_p, max_idx_l)
        accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

        saver = tf.train.Saver()
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            step = 0
            while True:
                batch_x, batch_y = self.get_next_batch(64)
                _, loss_ = sess.run([optimizer, loss],
                                    feed_dict={self.X: batch_x, self.Y: batch_y, self.keep_prob: 0.75,
                                               self.train_phase: True})
                print('step:%d,loss:%.5f' % (step, loss_))

                # 每100 step计算一次准确率
                if step % 100 == 0 and step != 0:
                    batch_x_test, batch_y_test = self.get_next_batch(100)
                    acc = sess.run(accuracy,
                                   feed_dict={self.X: batch_x_test, self.Y: batch_y_test, self.keep_prob: 1.,
                                              self.train_phase: False})
                    print("第%s步，训练准确率为：%s" % (step, acc))
                    # 如果准确率大80%,保存模型,完成训练
                    if loss_ < 0.003:
                        saver.save(sess, "./crack_capcha.model", global_step=step)
                        break
                    saver.save(sess, "./crack_capcha.model", global_step=step)
                step += 1

    def crack_captcha(self, img):
        # self.sess = tf.Session()
        # self.output = self.crack_captcha_cnn()
        # self.predict = tf.argmax(tf.reshape(self.output, [-1, self.MAX_CAPTCHA, self.CHAR_SET_LEN]), 2)
        # self.saver = tf.train.Saver()
        # print(settings.MODEL_DIR)
        # self.saver.restore(self.sess, tf.train.latest_checkpoint(settings.MODEL_DIR))
        img = img.convert("L")
        (width, height) = img.size
        pixles = img.load()
        for w in range(width):  # 二值化
            for h in range(height):
                if pixles[w, h] > 224:
                    pixles[w, h] = 255
                else:
                    pixles[w, h] = 0
        image = np.array(img.crop([0, 0, 160, 60]))
        image = self.convert2gray(image)  # 生成一张新图
        image = image.flatten() / 255  # 将图片一维化
        text_list = self.sess.run(self.predict, feed_dict={self.X: [image], self.keep_prob: 1, self.train_phase: False})
        text = text_list[0].tolist()
        vector = np.zeros(self.MAX_CAPTCHA * self.CHAR_SET_LEN)
        i = 0
        for n in text:
            vector[i * self.CHAR_SET_LEN + n] = 1
            i += 1
        captcha = self.vec2text(vector)
        return captcha
